import logging
import random

import numpy as np

from torch.utils.data import Dataset
from torchvision import transforms

from utils.augmentation import Augmentation
from tools.dataset_tools import dataset_tools
# from datasets.preprocessor import Preprocessor

from tools.sentinel12_tools import sentinel12_tools
from tools.os_tools import os_tools
from tools.QXSLAB_tools import QXSLAB_tools

logger = logging.getLogger("match")


DEF_AUG_PROBS = {
        "fliplr": 0.5,
        "flipud": 0.5,
        "scale": 0,
        "scale_px": (1.0, 1.0),
        "translate": 0,
        "translate_perc": (0.0, 0.0),
        "rotate": 0,
        "rotate_angle": (-5, 5)
    }


class PatchDataset(Dataset):
    # default args of augment
    AUG_PROBS = DEF_AUG_PROBS


    def __init__(
        self, config=None, *, tools=None, name="", modal="OPT", **kwargs
    ) -> None:
        super(PatchDataset, PatchDataset).__init__(self)

        if config == None:
            return
        
        self.modal = modal

        #-- transform
        func = []
        
        if config.augment:
            if config.augment is True or len(config.augment) == 0:
                config.augment = self.AUG_PROBS.copy()

            self.augmentor = Augmentation(probs=config.augment)
        else:
            self.augmentor = None

        func.append(transforms.ToTensor())
        self.transforms = transforms.Compose(func)

        #-- dataset
        tools = dataset_tools if tools is None else tools
        self.dt = tools(config.base_dir, config.list_file)
        self.pl = self.dt.get_pair_list()

        logger.info(f"{name} Patch Dataset created with {self.__len__()} patchs")


    def __getitem__(self, index):
        #-- augmentor
        if self.augmentor:
            self.augmentor.refresh_random_state()

        img = self._load_data(index)

        if self.augmentor is not None:
            img = self.augmentor(img)

        img = img[:, :, 0]
        img = np.ascontiguousarray(img)
        img = (img - img.min())/(img.ptp())
        
        img_tensor = self.transforms(img).float()

        return img_tensor


    def __len__(self):
        return len(self.pl)


    def _normalize_scale(self, x, in_range=(0, 255)):
        # clip 限定到 x 到 in_range 范围中，超出的作边界处理
        return (x.clip(*in_range) - in_range[0])/(in_range[1] - in_range[0])


    def _load_data(self, index, **kwargs):
        '''
        index to transformed image patch if exists

        Return:
        ------
            patch
        '''
        modal_ids = 0 if self.modal == "OPT" else 1
        im = self.dt.get_patch(
            self.pl[index][modal_ids], transform=self._normalize_scale, modal=self.modal
        )

        return im 


class Sentinel12OPTPatchDataset(PatchDataset):


    def __init__(self, config=None):
        super(Sentinel12OPTPatchDataset, self).__init__(
            config, tools=sentinel12_tools, name="Sentinel12"
        )


class OSOPTPatchDataset(PatchDataset):


    def __init__(self, config=None):
        super(OSOPTPatchDataset, self).__init__(
            config, tools=os_tools, name="OS"
        )


class OSSARPatchDataset(PatchDataset):


    def __init__(self, config=None):
        super(OSSARPatchDataset, self).__init__(
            config, tools=os_tools, name="OS", modal="SAR"
        )


class QXSLABOPTPatchDataset(PatchDataset):


    def __init__(self, config=None):
        super(QXSLABOPTPatchDataset, self).__init__(
            config, tools=QXSLAB_tools, name="QXSLAB"
        )


class QXSLABSARPatchDataset(PatchDataset):


    def __init__(self, config=None):
        super(QXSLABSARPatchDataset, self).__init__(
            config, tools=QXSLAB_tools, name="QXSLAB", modal="SAR"
        )


class NoisyPatchDataset(PatchDataset):


    def __init__(self, config=None, *, tools=None, name="", **kwargs) -> None:
        name = "Noisy" + name
        self.fn = self._get_noise_img
        self.scale = []
        for i in config.L:
            self.scale.append(1/i)

        super().__init__(config, tools=tools, name=name)


    def __getitem__(self, index):
        #-- augmentor
        if self.augmentor:
            self.augmentor.refresh_random_state()

        img = self._load_data(index)

        if self.augmentor is not None:
            img = self.augmentor(img)

        img = img[:, :, 0]
        img = np.ascontiguousarray(img)
        img = (img - img.min())/(img.ptp())

        img_noise = self.fn(img)

        img_tensor = self.transforms(img).float()
        noisy_img_tensor = self.transforms(img_noise).float()

        return img_tensor, noisy_img_tensor


    def _get_noise_img(self, img):
        # scale = np.random.rand()*.4
        # scale = 0.05 if scale < 0.05 else scale
        if len(self.scale) == 1:
            scale = self.scale[0]
        else:
            scale = random.choice(self.scale)

        noise = np.random.gamma(1/scale, scale, img.shape)

        img_noise = noise * img
        # img_noise = (img_noise-img_noise.min())/img_noise.ptp()
        img_noise = np.clip(img_noise, 0.0, 1.0)

        return img_noise


class Sentinel12OPTNoisyPatchDataset(NoisyPatchDataset):


    def __init__(self, config=None):
        super(Sentinel12OPTNoisyPatchDataset, self).__init__(
            config, tools=sentinel12_tools, name="Sentinel12"
        )
        

class OSOPTNoisyPatchDataset(NoisyPatchDataset):


    def __init__(self, config=None):
        super(OSOPTNoisyPatchDataset, self).__init__(
            config, tools=os_tools, name="OS"
        )
        

class QXSLABOPTNoisyPatchDataset(NoisyPatchDataset):


    def __init__(self, config=None):
        super(QXSLABOPTNoisyPatchDataset, self).__init__(
            config, tools=QXSLAB_tools, name="QXSLAB"
        )


if __name__ == "__main__":
    import cv2 as cv

    class config:
        def __init__(self):
            self.base_dir = None
            # self.list_file = "E:/workspace/SOMatch/json/sen12_train_val/train.json"
            # self.list_file = "E:/workspace/SOMatch/json/os_dataset/256_test.json"
            # self.list_file = "E:/workspace/SOMatch/json/QXSLAB_SAROPT/QingDao.json"
            self.list_file = "/home/rslab/DLMatch/workspace/SOMatch/json/os_dataset/512_test.json"
            self.augment = True


    cfg = config()

    # dset = Sentinel12OPTPatchDataset(cfg)
    # dset = OSOPTPatchDataset(cfg)
    # dset = QXSLABOPTPatchDataset(cfg)
    # dset = Sentinel12OPTNoisyPatchDataset(cfg)
    # dset = OSOPTNoisyPatchDataset(cfg)
    # dset = QXSLABOPTNoisyPatchDataset(cfg)
    dset = OSSARPatchDataset(cfg)

    img = dset[10]
    # img, noisy_img = dset[150]

    img = img.numpy()
    img = np.transpose(img, [1, 2, 0])

    # noisy_img = noisy_img.numpy()
    # noisy_img = np.transpose(noisy_img, [1, 2, 0])

    cv.imwrite("./tmp/img.png", img*255)
    # cv.imwrite("image/single_patch/noisy_img.png", noisy_img*255)

if __name__ == "__main__1":
    import os
    import cv2 as cv

    folder = "/home/rslab/DLMatch/datasets/despeckle"
    fd_kw = ["clean", "noisy"]

    L_list = [1, 2, 4, 8]

    class config:
        def __init__(self, l):
            self.base_dir = None
            self.list_file = "/home/rslab/DLMatch/workspace/SOMatch/json/QXSLAB_Despeckle/val39-1.json"
            self.augment = False
            self.L = [l]

    for L in L_list:

        cfg = config(L)
        dset = QXSLABOPTNoisyPatchDataset(cfg)

        cur_folder = os.path.join(folder, str(L))
        if not os.path.exists(cur_folder):
            os.makedirs(cur_folder)
            os.mkdir(os.path.join(cur_folder, "clean"))
            os.mkdir(os.path.join(cur_folder, "noisy"))

        clean_folder = os.path.join(cur_folder, "clean")
        noisy_folder = os.path.join(cur_folder, "noisy")

        cnt = 1
        for img, noisy_img in dset:

            img = img.numpy()
            img = np.transpose(img, [1, 2, 0])
            noisy_img = noisy_img.numpy()
            noisy_img = np.transpose(noisy_img, [1, 2, 0])
            
            filename = f"{cnt}.png"

            cv.imwrite(
                os.path.join(clean_folder, filename), 
                img*255)
            cv.imwrite(
                os.path.join(noisy_folder, filename), 
                noisy_img*255)

            cnt += 1
